#include "kalman_filter/kalman.h"

/*
 * filter: EKF filter structure
 */
void kf_predict(struct kalman_filter *filter, float *input, float dt)
{
	float AP_buf[KF_MAX_STATE_DIM * KF_MAX_STATE_DIM] = {0.0f};
	float trA_buf[KF_MAX_STATE_DIM * KF_MAX_STATE_DIM] = {0.0f};
	float APtrA_buf[KF_MAX_STATE_DIM * KF_MAX_STATE_DIM] = {0.0f};

	arm_matrix_instance_f32 AMat = {
		.numRows = filter->state_dim,
		.numCols = filter->state_dim,
		.pData = filter->A
	};
	arm_matrix_instance_f32 PMat = {
		.numRows = filter->state_dim,
		.numCols = filter->state_dim,
		.pData = filter->P
	};
	arm_matrix_instance_f32 QMat = {
		.numRows = filter->state_dim,
		.numCols = filter->state_dim,
		.pData = filter->Q
	};
	arm_matrix_instance_f32 APMat = {
		.numRows = filter->state_dim,
		.numCols = filter->state_dim,
		.pData = AP_buf
	};
	arm_matrix_instance_f32 trAMat = {
		.numRows = filter->state_dim,
		.numCols = filter->state_dim,
		.pData = trA_buf
	};
	arm_matrix_instance_f32 APtrAMat = {
		.numRows = filter->state_dim,
		.numCols = filter->state_dim,
		.pData = APtrA_buf
	};

	/* must be guaranteed */
	if (filter->x == NULL || filter->P == NULL || filter->A == NULL
		|| filter->Q == NULL || filter->process == NULL
		|| filter->computeA == NULL || filter->computeQ == NULL)
		return;

	/* must be guaranteed */
	if (filter->x == NULL || filter->P == NULL || filter->A == NULL
		|| filter->Q == NULL || filter->process == NULL
		|| filter->computeA == NULL || filter->computeQ == NULL)
		return;

	filter->u  =  input;
	filter->dt  =  dt;

	filter->computeA(filter);
	filter->computeQ(filter);
	filter->process(filter);

	arm_mat_trans_f32(&AMat, &trAMat);
	arm_mat_mult_f32(&AMat, &PMat, &APMat);
	arm_mat_mult_f32(&APMat, &trAMat, &APtrAMat);

	arm_mat_add_f32(&APtrAMat, &QMat, &PMat);
}

static inline void compute_K(struct kalman_filter *filter, float *K_buf)
{
	float trH_buf[KF_MAX_STATE_DIM * KF_MAX_MEASURE_DIM] = {0.0f};
	float PtrH_buf[KF_MAX_STATE_DIM * KF_MAX_MEASURE_DIM] = {0.0f};
	float HPtrH_buf[KF_MAX_MEASURE_DIM * KF_MAX_MEASURE_DIM] = {0.0f};
	float HPtrH_R_buf[KF_MAX_MEASURE_DIM * KF_MAX_MEASURE_DIM] = {0.0f};
	float invHPtrH_R_buf[KF_MAX_MEASURE_DIM * KF_MAX_MEASURE_DIM] = {0.0f};

	if (K_buf == NULL)
		return;

	arm_matrix_instance_f32 HMat = {.numRows = filter->measure_dim,
					.numCols = filter->state_dim,
					.pData = filter->H};
	arm_matrix_instance_f32 RMat = {.numRows = filter->measure_dim,
					.numCols = filter->measure_dim,
					.pData = filter->R};
	arm_matrix_instance_f32 PMat = {.numRows = filter->state_dim,
					.numCols = filter->state_dim,
					.pData = filter->P};
	arm_matrix_instance_f32 trHMat = {.numRows = filter->state_dim,
					.numCols = filter->measure_dim,
					.pData = trH_buf};
	arm_matrix_instance_f32 PtrHMat = {.numRows = filter->state_dim,
					.numCols = filter->measure_dim,
					.pData = PtrH_buf};
	arm_matrix_instance_f32 HPtrHMat = {.numRows = filter->measure_dim,
					.numCols = filter->measure_dim,
					.pData = HPtrH_buf};
	arm_matrix_instance_f32 HPtrH_RMat = {.numRows = filter->measure_dim,
					.numCols = filter->measure_dim,
					.pData = HPtrH_R_buf};
	arm_matrix_instance_f32 invHPtrH_RMat = {.numRows = filter->measure_dim,
					.numCols = filter->measure_dim,
					.pData = invHPtrH_R_buf};
	arm_matrix_instance_f32 KMat = {.numRows = filter->state_dim,
					.numCols = filter->measure_dim,
					.pData = K_buf};

	arm_mat_trans_f32(&HMat, &trHMat);
	arm_mat_mult_f32(&PMat, &trHMat, &PtrHMat);
	arm_mat_mult_f32(&HMat, &PtrHMat, &HPtrHMat);
	arm_mat_add_f32(&HPtrHMat, &RMat, &HPtrH_RMat);
	arm_mat_inverse_f32(&HPtrH_RMat, &invHPtrH_RMat);
	arm_mat_mult_f32(&PtrHMat, &invHPtrH_RMat, &KMat);	/* K = PH'/(HPH'+R) */
}

void kf_update(struct kalman_filter *filter,float *measure)
{
	uint8_t i;
	float *psrc;
	uint16_t ss_size = (uint16_t)(filter->state_dim)
			* (uint16_t)(filter->state_dim);

	float K_buf[KF_MAX_STATE_DIM * KF_MAX_MEASURE_DIM] = {0.0f};
	float esti_z_buf[KF_MAX_MEASURE_DIM] = {0.0f};
	float r_buf[KF_MAX_MEASURE_DIM] = {0.0f};
	float Kr_buf[KF_MAX_STATE_DIM] = {0.0f};
	float newx_buf[KF_MAX_STATE_DIM] = {0.0f};
	float KH_buf[KF_MAX_STATE_DIM * KF_MAX_STATE_DIM] = {0.0f};
	float I_buf[KF_MAX_STATE_DIM * KF_MAX_STATE_DIM] = {0.0f};
	float I_KH_buf[KF_MAX_STATE_DIM * KF_MAX_STATE_DIM] = {0.0f};
	float newP_buf[KF_MAX_STATE_DIM * KF_MAX_STATE_DIM] = {0.0f};

	arm_matrix_instance_f32 xMat = {.numRows = filter->state_dim,
					.numCols = 1,
					.pData = filter->x};
	arm_matrix_instance_f32 PMat = {.numRows = filter->state_dim,
					.numCols = filter->state_dim,
					.pData = filter->P};
	arm_matrix_instance_f32 HMat = {.numRows = filter->measure_dim,
					.numCols = filter->state_dim,
					.pData = filter->H};
	arm_matrix_instance_f32 KMat = {.numRows = filter->state_dim,
					.numCols = filter->measure_dim,
					.pData = K_buf};
	arm_matrix_instance_f32 esti_zMat = {.numRows = filter->measure_dim,
					.numCols = 1,
					.pData = esti_z_buf};
	arm_matrix_instance_f32 zMat = {.numRows = filter->measure_dim,
					.numCols = 1,
					.pData = measure};
	arm_matrix_instance_f32 rMat = {.numRows = filter->measure_dim,
					.numCols = 1,
					.pData = r_buf};
	arm_matrix_instance_f32 KrMat = {.numRows = filter->state_dim,
					.numCols = 1,
					.pData = Kr_buf};
	arm_matrix_instance_f32 newxMat = {.numRows = filter->state_dim,
					.numCols = 1,
					.pData = newx_buf};
	arm_matrix_instance_f32 KHMat = {.numRows = filter->state_dim,
					.numCols = filter->state_dim,
					.pData = KH_buf};
	arm_matrix_instance_f32 IMat = {.numRows = filter->state_dim,
					.numCols = filter->state_dim,
					.pData = I_buf};
	arm_matrix_instance_f32 I_KHMat = {.numRows = filter->state_dim,
					.numCols = filter->state_dim,
					.pData = I_KH_buf};
	arm_matrix_instance_f32 newPMat = {.numRows = filter->state_dim,
					.numCols = filter->state_dim,
					.pData = newP_buf};

	if (filter->x == NULL || filter->P == NULL || filter->H == NULL
		|| filter->R == NULL || filter->measurement == NULL
		|| filter->computeH == NULL || filter->computeR == NULL)
		return;

	if (filter->x == NULL || filter->P == NULL || filter->H == NULL
		|| filter->R == NULL || filter->measurement == NULL
		|| filter->computeH == NULL || filter->computeR == NULL)
		return;

	filter->z = measure;

	filter->computeH(filter);
	filter->computeR(filter);

	compute_K(filter, K_buf);	/* K = PH'/(HPH'+R) */

	filter->measurement((void*)filter, esti_z_buf);		/* z_tilde = h(x) */
	arm_mat_sub_f32(&zMat, &esti_zMat, &rMat);		/* r = z - z_tilde */
	arm_mat_mult_f32(&KMat, &rMat, &KrMat);
	arm_mat_add_f32(&xMat, &KrMat, &newxMat);		/* x = x + Kr */

	arm_copy_f32(newxMat.pData, filter->x, filter->state_dim);

	arm_mat_mult_f32(&KMat, &HMat, &KHMat);

	for (i = 0, psrc = IMat.pData;
	     i < filter->state_dim;
	     i++, psrc += (filter->state_dim + 1)) {
		*psrc = 1.0f;
	}

	arm_mat_sub_f32(&IMat, &KHMat, &I_KHMat);
	arm_mat_mult_f32(&I_KHMat, &PMat, &newPMat);	/* P = (I - KH)P */

	arm_copy_f32(newPMat.pData, filter->P, ss_size);
}
